# -*- coding:utf-8 -*-

import torch

#预训练模型来源
source = 'huggingface/pytorch-transformers'

#加载的预训练模型的名字
model_name = 'bert-base-chinese'

#选定加载模型的哪⼀部分,这⾥是模型的映射器
part = 'tokenizer'
tokenizer = torch.hub.load(source, part, model_name)

# # 加载不带头的预训练模型
part = 'model'
model = torch.hub.load(source, part, model_name)

# # 加载带有语⾔模型头的预训练模型
# part = 'modelWithLMHead'
# lm_model = torch.hub.load(source, part, model_name)

# 加载带有分类模型头的预训练模型
part = 'modelForSequenceClassification'
classification_model = torch.hub.load(source, part, model_name)

# 加载带有问答模型头的预训练模型
part = 'modelForQuestionAnswering'
qa_model = torch.hub.load(source, part, model_name)

##################################################################################################

# 输⼊的中⽂⽂本
input_text = "人生该如何起头"

# 使⽤tokenizer进⾏数值映射
indexed_tokens = tokenizer.encode(input_text)

# 打印映射后的结构
print("indexed_tokens:", indexed_tokens)

# 将映射结构转化为张量输送给不带头的预训练模型
tokens_tensor = torch.tensor([indexed_tokens])

# 使⽤不带头的预训练模型获得结果
with torch.no_grad():
    model_result = model(tokens_tensor)
    encoded_layers = model_result.last_hidden_state
print("不带头的模型输出结果:", encoded_layers)
print("不带头的模型输出结果的尺⼨:", encoded_layers.shape)

# 使⽤带有分类模型头的预训练模型获得结果
with torch.no_grad():
    classification_model_result = classification_model(tokens_tensor)
    classification_output = classification_model_result.logits
print("带有分类模型头的预训练模型输出结果:", classification_output)
print("带有分类模型头的预训练模型输出结果的尺⼨:", classification_output.shape)


input_text1 = "我家的小狗是黑色的"
input_text2 = "我家的小狗是什么颜色的呢？"

#映射两个句子
indexed_tokens = tokenizer.encode(input_text1, input_text2)
print("句子对的indexed_tokens:", indexed_tokens)

# 用0, 1来区分第一条和第二条句子
segments_ids = [0]*(len(input_text1)+2) + [1]*(len(input_text2)+1)

# 转化张量形式
segments_tensors = torch.tensor([segments_ids])
tokens_tensor = torch.tensor([indexed_tokens])

# 使用带有问答模型头的预训练模型获得结果
with (torch.no_grad()):
    qa_model_result = qa_model(tokens_tensor, token_type_ids=segments_tensors)
    start_logits = qa_model_result.start_logits
    end_logits = qa_model_result.end_logits

print("带问答模型的模型输出结果：", (start_logits, end_logits))
print("带问答模型的模型输出结果的尺寸：", (start_logits.shape, end_logits.shape))